#!/usr/bin/env python
# encoding: utf-8
'''
@author: wangjianrong
@software: PyCharm
@file: resize_model.py
@time: 2019/9/2 13:26
@desc: 对coco上训练的模型进行修改，根据不同的num_classes和num_anchors修改模型参数size，使finetune能够能够正常加载与训练模型
        该文件针对cascade_rcnn进行修改
'''

import torch
import torch.nn as nn
import os

check_file = 'weights/gfl_x101_32x4d_fpn_dconv_c4-c5_mstrain_2x_coco_20200630_102002-14a2bf25.pth'
check_name = os.path.basename(check_file)
check_name = check_name[:check_name.rfind('.')]
num_classes = 1 #不包括背景
pre_weights = torch.load(check_file)['state_dict']
for name,param in pre_weights.items():
    print(name,param.shape)
pre_weights['bbox_head.gfl_cls.weight'].resize_(num_classes,256,3,3)
pre_weights['bbox_head.gfl_cls.bias'].resize_(num_classes)

#此处需要重新初始化，resize后赋值通常会很大，导致训练梯度爆炸，loss inf
nn.init.kaiming_normal_(pre_weights['bbox_head.gfl_cls.weight'], a=0, mode='fan_out', nonlinearity='relu')
nn.init.constant_(pre_weights['bbox_head.gfl_cls.bias'], 0)


torch.save(pre_weights,'weights/{}_{}cls.pth'.format(check_name,num_classes))